package com.gitee.helioz.openai.api;

import com.gitee.helioz.openai.common.error.OpenAiRuntimeException;
import com.gitee.helioz.openai.common.util.http.ContentTypeEnum;
import com.gitee.helioz.openai.common.util.http.HttpFilePart;
import com.gitee.helioz.openai.common.util.http.HttpProxyConfig;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Authenticator;
import okhttp3.Call;
import okhttp3.ConnectionPool;
import okhttp3.Credentials;
import okhttp3.Dispatcher;
import okhttp3.FormBody;
import okhttp3.MediaType;
import okhttp3.MultipartBody;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;
import okhttp3.ResponseBody;
import org.apache.commons.lang3.StringUtils;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Proxy;
import java.net.SocketAddress;
import java.util.List;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * 〈OkHttp实现〉
 *
 * @author <a href="https://gitee.com/helioz">Morning Starr</a>
 */
@Slf4j
public class OpenAiOkHttpServiceImpl extends BaseOpenAiServiceImpl {

    private OkHttpClient httpClient;
    private OkHttpClient.Builder httpClientBuilder;

    public OpenAiOkHttpServiceImpl() {
        super();
    }

    public OpenAiOkHttpServiceImpl(HttpProxyConfig httpProxyConfig) {
        super(httpProxyConfig);
    }

    public OpenAiOkHttpServiceImpl(OkHttpClient.Builder httpClientBuilder) {
        this.httpClientBuilder = httpClientBuilder;
    }

    public OpenAiOkHttpServiceImpl(OkHttpClient.Builder httpClientBuilder, HttpProxyConfig httpProxyConfig) {
        super(httpProxyConfig);
        this.httpClientBuilder = httpClientBuilder;
    }

    @Override
    protected void initHttp() {
        if (this.httpClientBuilder == null) {
            httpClientBuilder = new OkHttpClient.Builder();
            // 连接池
            ConnectionPool connectionPool = new ConnectionPool(50, 30, TimeUnit.SECONDS);
            httpClientBuilder.connectionPool(connectionPool);
            // 超时时间
            httpClientBuilder.connectTimeout(10, TimeUnit.SECONDS)
                    .readTimeout(10, TimeUnit.SECONDS)
                    .writeTimeout(10, TimeUnit.SECONDS)
                    .connectionPool(connectionPool)
                    .retryOnConnectionFailure(false);
            // 线程池管理
            int cpuCount = Runtime.getRuntime().availableProcessors();
            int corePoolSize = cpuCount + 1;
            int maximumPoolSize = corePoolSize;
            long keepAliveTime = 60L;
            TimeUnit unit = TimeUnit.SECONDS;
            BlockingQueue<Runnable> workQueue = new LinkedBlockingQueue<>(1000);
            ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor(
                    corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, r -> {
                Thread t = new Thread(r);
                t.setName("OKHTTP-THREAD-POOL-" + t.getId());
                return t;
            }
            );
            httpClientBuilder.dispatcher(new Dispatcher(threadPoolExecutor));
            // 设置代理
            if (this.getHttpProxyConfig() != null) {
                HttpProxyConfig httpProxyConfig = this.getHttpProxyConfig();
                SocketAddress proxyAddress = new InetSocketAddress(httpProxyConfig.getHost(), httpProxyConfig.getPort());
                Proxy proxy = new Proxy(httpProxyConfig.getType(), proxyAddress);
                httpClientBuilder.proxy(proxy);
                // 设置代理认证账号密码
                if (!StringUtils.isAnyBlank(httpProxyConfig.getName(), httpProxyConfig.getPassword())) {
                    Authenticator proxyAuthenticator = (route, response) -> {
                        String credential = Credentials.basic(httpProxyConfig.getName(), httpProxyConfig.getPassword());
                        return response.request().newBuilder()
                                .header("Proxy-Authorization", credential)
                                .build();
                    };
                    httpClientBuilder.authenticator(proxyAuthenticator);
                }
            }
        }
        // 创建HttpClient对象
        this.httpClient = this.httpClientBuilder.build();
        log.info("初始化OkHttp Client4 完成");
    }

    @Override
    protected String sendGetRequest(String url) throws IOException {
        Request.Builder requestBuilder = new Request.Builder();
        requestBuilder.url(url);
        Map<String, String> headerMap = buildHeaderMap();
        headerMap.forEach(requestBuilder::addHeader);
        Request request = requestBuilder.build();
        Call call = httpClient.newCall(request);
        Response response = call.execute();
        ResponseBody responseBody = response.body();
        if (responseBody == null) {
            throw new OpenAiRuntimeException("响应体为空");
        }
        return responseBody.string();
    }

    @Override
    protected String sendPostJsonRequest(String url, String jsonBody) throws IOException {
        Request.Builder requestBuilder = new Request.Builder();
        requestBuilder.url(url);
        Map<String, String> headerMap = buildHeaderMap();
        headerMap.forEach(requestBuilder::addHeader);
        RequestBody requestBody = RequestBody.create(jsonBody, MediaType.parse(ContentTypeEnum.JSON.getType()));
        requestBuilder.post(requestBody);
        Request request = requestBuilder.build();
        Call call = httpClient.newCall(request);
        Response response = call.execute();
        ResponseBody responseBody = response.body();
        if (responseBody == null) {
            throw new OpenAiRuntimeException("响应体为空");
        }
        return responseBody.string();
    }

    @Override
    protected String sendPostFormData(String url, Map<String, Object> formMap, List<HttpFilePart> fileParts) throws IOException {
        Request.Builder requestBuilder = new Request.Builder();
        requestBuilder.url(url);
        Map<String, String> headerMap = buildHeaderMap();
        headerMap.forEach(requestBuilder::addHeader);
        RequestBody requestBody;
        if (fileParts != null && !fileParts.isEmpty()) {
            MultipartBody.Builder mpBuilder = new MultipartBody.Builder();
            mpBuilder.setType(MultipartBody.FORM);
            for (Map.Entry<String, Object> entry : formMap.entrySet()) {
                if (entry.getValue() != null) {
                    mpBuilder.addFormDataPart(entry.getKey(), String.valueOf(entry.getValue()));
                }
            }
            for (HttpFilePart filePart : fileParts) {
                mpBuilder.addFormDataPart(filePart.getPartName(), filePart.getFileName(), RequestBody.create(filePart.getFileBytes()));
            }
            requestBody = mpBuilder.build();
        } else {
            FormBody.Builder formBuilder = new FormBody.Builder();
            for (Map.Entry<String, Object> entry : formMap.entrySet()) {
                if (entry.getValue() != null) {
                    formBuilder.add(entry.getKey(), String.valueOf(entry.getValue()));
                }
            }
            requestBody = formBuilder.build();
        }
        Request request = requestBuilder.post(requestBody).build();
        Call call = httpClient.newCall(request);
        Response response = call.execute();
        ResponseBody responseBody = response.body();
        if (responseBody == null) {
            throw new OpenAiRuntimeException("响应体为空");
        }
        return responseBody.string();
    }

    @Override
    protected String sendDeleteRequest(String url) throws IOException {
        Request.Builder requestBuilder = new Request.Builder();
        requestBuilder.url(url);
        requestBuilder.delete();
        Map<String, String> headerMap = buildHeaderMap();
        headerMap.forEach(requestBuilder::addHeader);
        Request request = requestBuilder.build();
        Call call = httpClient.newCall(request);
        Response response = call.execute();
        ResponseBody responseBody = response.body();
        if (responseBody == null) {
            throw new OpenAiRuntimeException("响应体为空");
        }
        return responseBody.string();
    }
}
